// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <kllm/utility/cpu.h>
#define JSON_ASSERT GGML_ASSERT
#include <nlohmann/json.hpp>
#include <kllm/utility/cpu.h>
#include <algorithm>
#include <cinttypes>
#include <climits>
#include <cmath>
#include <codecvt>
#include <cstdarg>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <iterator>
#include <regex>
#include <sstream>
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <turbo/log/logging.h>

#if defined(__APPLE__) && defined(__MACH__)
#include <sys/types.h>
#include <sys/sysctl.h>
#include <sys/types.h>
#include <sys/resource.h>
#endif

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#ifndef NOMINMAX
#   define NOMINMAX
#endif
#include <locale>
#include <windows.h>
#include <fcntl.h>
#include <io.h>
#else
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/resource.h>
#endif
#if defined(LLAMA_USE_CURL)
#include <curl/curl.h>
#include <curl/easy.h>
#include <future>
#endif

#if defined(_MSC_VER)
#pragma warning(disable: 4244 4267) // possible loss of data
#endif

namespace kllm {


    using json = nlohmann::ordered_json;

//
// CPU utils
//

    int32_t cpu_get_num_physical_cores() {
#ifdef __linux__
        // enumerate the set of thread siblings, num entries is num cores
        std::unordered_set<std::string> siblings;
        for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) {
            std::ifstream thread_siblings("/sys/devices/system/cpu/cpu"
                                          + std::to_string(cpu) + "/topology/thread_siblings");
            if (!thread_siblings.is_open()) {
                break; // no more cpus
            }
            std::string line;
            if (std::getline(thread_siblings, line)) {
                siblings.insert(line);
            }
        }
        if (!siblings.empty()) {
            return static_cast<int32_t>(siblings.size());
        }
#elif defined(__APPLE__) && defined(__MACH__)
        int32_t num_physical_cores;
    size_t len = sizeof(num_physical_cores);
    int result = sysctlbyname("hw.perflevel0.physicalcpu", &num_physical_cores, &len, NULL, 0);
    if (result == 0) {
        return num_physical_cores;
    }
    result = sysctlbyname("hw.physicalcpu", &num_physical_cores, &len, NULL, 0);
    if (result == 0) {
        return num_physical_cores;
    }
#elif defined(_WIN32) && (_WIN32_WINNT >= 0x0601) && !defined(__MINGW64__) // windows 7 and later
    // TODO: windows + arm64 + mingw64
    unsigned int n_threads_win = std::thread::hardware_concurrency();
    unsigned int default_threads = n_threads_win > 0 ? (n_threads_win <= 4 ? n_threads_win : n_threads_win / 2) : 4;

    DWORD buffer_size = 0;
    if (!GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &buffer_size)) {
        if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
            return default_threads;
        }
    }

    std::vector<char> buffer(buffer_size);
    if (!GetLogicalProcessorInformationEx(RelationProcessorCore, reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(buffer.data()), &buffer_size)) {
        return default_threads;
    }

    int32_t num_physical_cores = 0;
    PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(buffer.data());
    while (buffer_size > 0) {
        if (info->Relationship == RelationProcessorCore) {
            num_physical_cores += info->Processor.GroupCount;
        }
        buffer_size -= info->Size;
        info = reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(reinterpret_cast<char*>(info) + info->Size);
    }

    return num_physical_cores > 0 ? num_physical_cores : default_threads;
#endif
        unsigned int n_threads = std::thread::hardware_concurrency();
        return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
    }

#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__)
#include <pthread.h>

static void cpuid(unsigned leaf, unsigned subleaf,
                  unsigned *eax, unsigned *ebx, unsigned *ecx, unsigned *edx) {
    __asm__("movq\t%%rbx,%%rsi\n\t"
            "cpuid\n\t"
            "xchgq\t%%rbx,%%rsi"
            : "=a"(*eax), "=S"(*ebx), "=c"(*ecx), "=d"(*edx)
            : "0"(leaf), "2"(subleaf));
}

static int pin_cpu(int cpu) {
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(cpu, &mask);
    return pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask);
}

static bool is_hybrid_cpu(void) {
    unsigned eax, ebx, ecx, edx;
    cpuid(7, 0, &eax, &ebx, &ecx, &edx);
    return !!(edx & (1u << 15));
}

static bool is_running_on_efficiency_core(void) {
    unsigned eax, ebx, ecx, edx;
    cpuid(0x1a, 0, &eax, &ebx, &ecx, &edx);
    int intel_atom = 0x20;
    int core_type = (eax & 0xff000000u) >> 24;
    return core_type == intel_atom;
}

static int cpu_count_math_cpus(int n_cpu) {
    int result = 0;
    for (int cpu = 0; cpu < n_cpu; ++cpu) {
        if (pin_cpu(cpu)) {
            return -1;
        }
        if (is_running_on_efficiency_core()) {
            continue; // efficiency cores harm lockstep threading
        }
        ++cpu; // hyperthreading isn't useful for linear algebra
        ++result;
    }
    return result;
}

#endif // __x86_64__ && __linux__

/**
 * Returns number of CPUs on system that are useful for math.
 */
int32_t cpu_get_num_math() {
#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__)
    int n_cpu = sysconf(_SC_NPROCESSORS_ONLN);
    if (n_cpu < 1) {
        return cpu_get_num_physical_cores();
    }
    if (is_hybrid_cpu()) {
        cpu_set_t affinity;
        if (!pthread_getaffinity_np(pthread_self(), sizeof(affinity), &affinity)) {
            int result = cpu_count_math_cpus(n_cpu);
            pthread_setaffinity_np(pthread_self(), sizeof(affinity), &affinity);
            if (result > 0) {
                return result;
            }
        }
    }
#endif
    return cpu_get_num_physical_cores();
}

// Helper for setting process priority

#if defined(_WIN32)

bool set_process_priority(enum ggml_sched_priority prio) {
    if (prio == GGML_SCHED_PRIO_NORMAL) {
        return true;
    }

    DWORD p = NORMAL_PRIORITY_CLASS;
    switch (prio) {
        case GGML_SCHED_PRIO_NORMAL:   p = NORMAL_PRIORITY_CLASS;       break;
        case GGML_SCHED_PRIO_MEDIUM:   p = ABOVE_NORMAL_PRIORITY_CLASS; break;
        case GGML_SCHED_PRIO_HIGH:     p = HIGH_PRIORITY_CLASS;         break;
        case GGML_SCHED_PRIO_REALTIME: p = REALTIME_PRIORITY_CLASS;     break;
    }

    if (!SetPriorityClass(GetCurrentProcess(), p)) {
        LOG_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError());
        return false;
    }

    return true;
}

#else // MacOS and POSIX
    bool set_process_priority(enum ggml_sched_priority prio) {
        if (prio == GGML_SCHED_PRIO_NORMAL) {
            return true;
        }

        int p = 0;
        switch (prio) {
            case GGML_SCHED_PRIO_NORMAL:   p =  0;  break;
            case GGML_SCHED_PRIO_MEDIUM:   p = -5;  break;
            case GGML_SCHED_PRIO_HIGH:     p = -10; break;
            case GGML_SCHED_PRIO_REALTIME: p = -20; break;
        }

        if (!setpriority(PRIO_PROCESS, 0, p)) {
            PLOG(WARNING)<<"failed to set process priority "<<prio;
            return false;
        }
        return true;
    }

#endif

//
// CLI argument parsing
//


    void postprocess_cpu_params(cpu_params& cpuparams, const cpu_params* role_model) {
        int32_t n_set = 0;

        if (cpuparams.n_threads < 0) {
            // Assuming everything about cpuparams is invalid
            if (role_model != nullptr) {
                cpuparams = *role_model;
            } else {
                cpuparams.n_threads = cpu_get_num_math();
            }
        }

        for (int32_t i = 0; i < GGML_MAX_N_THREADS; i++) {
            if (cpuparams.cpumask[i]) {
                n_set++;
            }
        }

        if (n_set && n_set < cpuparams.n_threads) {
            // Not enough set bits, may experience performance issues.
            LOG(WARNING)<<"Not enough set bits in CPU mask ("<<n_set<<") to satisfy requested thread count: "<<cpuparams.n_threads;
        }
    }

    bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THREADS]) {
        size_t dash_loc = range.find('-');
        if (dash_loc == std::string::npos) {
            LOG(ERROR)<<"Format of CPU range is invalid! Expected [<start>]-[<end>].";
            return false;
        }

        size_t start_i;
        size_t end_i;

        if (dash_loc == 0) {
            start_i = 0;
        } else {
            start_i = std::stoull(range.substr(0, dash_loc));
            if (start_i >= GGML_MAX_N_THREADS) {
                LOG(ERROR)<<"Start index out of bounds!";
                return false;
            }
        }

        if (dash_loc == range.length() - 1) {
            end_i = GGML_MAX_N_THREADS - 1;
        } else {
            end_i = std::stoull(range.substr(dash_loc + 1));
            if (end_i >= GGML_MAX_N_THREADS) {
                LOG(ERROR)<<"End index out of bounds!\n";
                return false;
            }
        }

        for (size_t i = start_i; i <= end_i; i++) {
            boolmask[i] = true;
        }

        return true;
    }


    bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREADS]) {
        // Discard potential 0x prefix
        size_t start_i = 0;
        if (mask.length() >= 2 && mask.substr(0, 2) == "0x") {
            start_i = 2;
        }

        size_t num_digits = mask.length() - start_i;
        if (num_digits > 128) num_digits = 128;

        size_t end_i = num_digits + start_i;

        for (size_t i = start_i, n = (num_digits*4 - 1); i < end_i; i++, n-=4) {
            char c = mask.at(i);
            int8_t id = c;

            if ((c >= '0' && c <= '9')) {
                id -= '0';
            } else if (c >= 'a' && c <= 'f') {
                id -= 'a' - 10;
            } else if (c >= 'A' && c <= 'F') {
                id -= 'A' - 10;
            } else {
                LOG(ERROR)<<"Invalid hex character '"<<c<<"' at position "<<int32_t(i);
                return false;
            }

            boolmask[  n  ] = boolmask[  n  ] || ((id & 8) != 0);
            boolmask[n - 1] = boolmask[n - 1] || ((id & 4) != 0);
            boolmask[n - 2] = boolmask[n - 2] || ((id & 2) != 0);
            boolmask[n - 3] = boolmask[n - 3] || ((id & 1) != 0);
        }

        return true;
    }

}  // namespace kllm
